"""
Script containing `Pydantic <https://docs.pydantic.dev/latest/>`_ models for
YAML configuration files.

.. note::

   ``models.py`` does *not* depend on ``ng911ok``.

This script is intended to read, validate, and generate JSON schemas for all
configuration files included in the Toolkit. The :class:`TagFreeLoader` class
ignores all YAML tags, so the data types used in the Pydantic models correspond
to the types of the values in the actual YAML files before they are processed
by the Toolkit.

The model classes should have attributes corresponding to keys in the YAML
files, and the attributes should have type hints similar to::

    Annotated[int, Field(description="Some integer")]

In the above case, ``int`` is the attribute's type, and
``Field(description="Some integer")`` provides additional information to
Pydantic. The ``description`` argument is used to generate the values of
``description`` keys for each item of ``properties`` keys in the output JSON
schemas.
"""

import argparse
import json
from pathlib import Path
from typing import Annotated, Optional, Final

from pydantic import BaseModel, Field, ValidationError, ConfigDict, TypeAdapter
import yaml


PATH_TOOLKIT_ROOT: Final[Path] = Path(__file__).parent.parent.parent
"""The root folder of the Toolkit."""

PATH_SCHEMA_ROOT: Final[Path] = Path(__file__).parent
"""The folder to which JSON schema files should be written."""


class TagFreeLoader(yaml.SafeLoader):
    """YAML loader that ignores tags."""

    def construct_object(self, node, deep=False):
        """Attempts to construct an object from *node* as normal. If a
        ``yaml.constructor.ConstructorError`` occurs, attempts to construct
        *node* as if it has no tag."""
        try:
            return super().construct_object(node, deep=deep)
        except yaml.constructor.ConstructorError:
            if isinstance(node, yaml.MappingNode):
                return self.construct_mapping(node, deep=deep)
            elif isinstance(node, yaml.SequenceNode):
                return self.construct_sequence(node, deep=deep)
            elif isinstance(node, yaml.ScalarNode):
                return self.construct_scalar(node)
            else:
                return None


#### TOPOLOGY ####

class OneMemberRule(BaseModel):
    rule: Annotated[str, Field(description="Topology rule to enforce")]
    member: Annotated[str, Field(description="Single feature class role to which the rule applies")]


class TwoMemberRule(BaseModel):
    member1: Annotated[str, Field(description="Role of the first feature class in the relationship")]
    rule: Annotated[str, Field(description="Topology rule to enforce")]
    member2: Annotated[str, Field(description="Role of the second feature class in the relationship")]


class TopologyModel(BaseModel):
    exception_field: Annotated[str, Field(description="Role of the topology exception field")]
    exception_domain: Annotated[str, Field(description="Domain that applies to the topology exception field")]
    required_dataset_topology_name: Annotated[str, Field(description="Name of the topology dataset in the required feature dataset")]
    optional_dataset_topology_name: Annotated[str, Field(description="Name of the topology dataset in the optional feature dataset")]
    required_dataset_rules: Annotated[list[OneMemberRule | TwoMemberRule], Field(description="Topology rules to enforce")]


#### COUNTIES ####

class CountyModel(BaseModel):
    state: Annotated[str, Field(description="State containing the county")]
    name: Annotated[str, Field(description="Name of the county")]
    fips3: Annotated[str, Field(description="The three digits of the county's FIPS code indictating only the county, not the state")]
    number: Annotated[int | None, Field(description="Number assigned to the county by its state")]


CountiesRootModel = TypeAdapter(list[CountyModel])


#### LEGACY ####

class LegacyFieldModel(BaseModel):
    next_gen: Annotated[str, Field(description="Role of the Next-Generation field in the relationship")]
    legacy: Annotated[str, Field(description="Role of the Legacy field in the relationship")]
    concatenation: Annotated[bool, Field(description="Whether the fields in the relationship are concatenations of other fields")]
    equal: Annotated[bool, Field(description="Whether the values of the related fields should be equal")]
    value_map: Annotated[dict[str, str] | None, Field(None, description="Mapping of Next-Generation to Legacy values")]

class LegacyModel(BaseModel):
    model_config = ConfigDict(extra="allow")

    fields: Annotated[list[LegacyFieldModel], Field(description="Specifications of relationships between Next-Generation and Legacy fields")]


#### CONFIG ####

class GDBInfoModel(BaseModel):
    spatial_reference_factory_code_2d: Annotated[int, Field(description="Factory code for the spatial reference to be used for all 2D NG911 feature classes")]
    spatial_reference_factory_code_3d: Annotated[int, Field(description="Factory code for the spatial reference to be used for all 3D NG911 feature classes")]
    required_dataset_name: Annotated[str, Field(description="Name of the dataset containing required feature classes")]
    optional_dataset_name: Annotated[str, Field(description="Name of the dataset containing optional feature classes")]
    nguid_urn_prefix: Annotated[str, Field(title="NGUID URN Prefix", description="The prefix, not including the trailing colon, common to every NGUID URN")]


class DomainModel(BaseModel):
    name: Annotated[str, Field(description="Name of the domain")]
    description: Annotated[str, Field(description="Description of the domain")]
    type: Annotated[str, Field(description="Type of domain")]
    entries: Annotated[dict[str, str], Field(description="Mapping of domain codes to descriptions")]


class FieldModel(BaseModel):
    role: Annotated[str, Field(description="Role of the field which should not change when the Standards are updated")]
    name: Annotated[str, Field(description="Name of the field which may change when the Standards are updated")]
    type: Annotated[str, Field(description="Data type of the field")]
    priority: Annotated[str, Field(description="Priority of the field as specified in the Standards")]
    length: Annotated[int | None, Field(None, description="Field length")]
    domain: Annotated[DomainModel | None, Field(None, description="Domain that applies to the field")]
    fill_value: Annotated[int | float | str | None, Field(None, description="Value to use as a default")]


class FeatureClassModel(BaseModel):
    role: Annotated[str, Field(description="Role of the feature class which should not change when the Standards are updated")]
    name: Annotated[str, Field(description="Name of the feature class which may change when the Standards are updated")]
    geometry_type: Annotated[str, Field(description="Geometry type of the feature class")]
    dataset: Annotated[str, Field(description="The feature dataset to which the feature class belongs")]
    unique_id: Annotated[FieldModel, Field(description="NGUID field for the feature class")]
    fields: Annotated[list[FieldModel], Field(description="Fields that the feature class should have")]


class ConfigRootModel(BaseModel):
    gdb_info: Annotated[GDBInfoModel, Field(description="Information about the geodatabase and its components")]
    domains: Annotated[list[DomainModel], Field(description="Domains the geodatabase should contain")]
    fields: Annotated[list[FieldModel], Field(description="All fields that are found in any of the feature classes")]
    feature_classes: Annotated[list[FeatureClassModel], Field(description="All feature classes, required or optional, that may be present in the geodatabase")]


def validate_file(
        path: Path,
        model: type[BaseModel] | TypeAdapter,
        quiet: bool = True,
        out_schema_file: Optional[str] = None
) -> bool:
    """
    Validates a YAML configuration file against a Pydantic model. Optionally
    prints the result of the validation and exports a JSON schema file.

    :param path: Path to the YAML configuration file
    :param model: Subclass of ``BaseModel`` or instance of ``TypeAdapter``
        against which the data in *path* should be validated
    :param quiet: If True, prints nothing to the console; default False
    :param out_schema_file: If supplied (and the data passes validation), the
        name of the JSON schema file to create in the :data:`PATH_SCHEMA_ROOT``
        directory
    :return: Whether the contents of the file at *path* passed validation
    """
    with open(path, "r") as f:
        data = yaml.load(f, Loader=TagFreeLoader)  # type: ignore
    try:
        schema: dict
        if isinstance(model, TypeAdapter):
            model.validate_python(data)
            schema = model.json_schema()
        elif issubclass(model, BaseModel):
            model.model_validate(data)
            schema = model.model_json_schema()
        else:
            raise TypeError
    except ValidationError as exc:
        if not quiet:
            print(f"\033[30m\033[41mFAIL: {path}\033[0m")  # Codes like '\033[30m' are ANSI color/style codes for console output
            print(f"\033[0;31m{exc}\033[0m")
        return False
    else:
        if not quiet:
            print(f"\033[30m\033[42mPASS: {path}\033[0m")
        if out_schema_file:
            with open(PATH_SCHEMA_ROOT / out_schema_file, "w") as f:
                f.write(json.dumps(schema, indent=2))
            if not quiet:
                print(f"Wrote schema to '{PATH_SCHEMA_ROOT / out_schema_file}'.")
        return True


class _Args:
    emit: bool
    quiet: bool


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--no-emit", action="store_false", help="Do not write JSON schema files", dest="emit")
    parser.add_argument("--quiet", action="store_true", help="Do not print validation results")
    args = parser.parse_args(namespace=_Args())

    validate_file(PATH_TOOLKIT_ROOT / "topology.yml", TopologyModel, args.quiet, "topologyschema.json" if args.emit else None)
    validate_file(PATH_TOOLKIT_ROOT / "counties.yml", CountiesRootModel, args.quiet, "countiesschema.json" if args.emit else None)
    validate_file(PATH_TOOLKIT_ROOT / "legacy.yml", LegacyModel, args.quiet, "legacyschema.json" if args.emit else None)
    validate_file(PATH_TOOLKIT_ROOT / "config.yml", ConfigRootModel, args.quiet, "configschema.json" if args.emit else None)


__all__ = ["PATH_TOOLKIT_ROOT", "PATH_SCHEMA_ROOT", "TagFreeLoader", "validate_file"]


if __name__ == "__main__":
    main()
